# 导入必要库
import tensorflow as tf
from tensorflow.keras import layers, datasets
import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

plt.rcParams['font.family'] = 'sans-serif'
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Arial']
mpl.rcParams['axes.unicode_minus'] = False
# ==================== 步骤1：加载并预处理MNIST数据集 ====================
# MNIST是手写数字数据集（28x28灰度图，0-9共10类）
(x_train, y_train), (x_test, y_test) = datasets.mnist.load_data()

# 数据归一化（像素值从0-255缩放到0-1）
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# 调整输入形状（28x28 -> 784维向量）
x_train = x_train.reshape(-1, 28 * 28)  # -1表示自动计算样本数
x_test = x_test.reshape(-1, 28 * 28)

# ==================== 步骤2：构建神经网络模型 ====================
model = tf.keras.Sequential([
    layers.Input(shape=(784,)),  # 输入层（784维特征）
    layers.Dense(128, activation='relu'),  # 隐藏层（128个神经元，ReLU激活）
    layers.Dense(64, activation='relu'),   # 隐藏层（64个神经元，ReLU激活）
    layers.Dense(10, activation='softmax')  # 输出层（10个神经元，Softmax激活，对应0-9类别概率）
])

# 查看模型结构
model.summary()

# ==================== 步骤3：编译模型 ====================
model.compile(
    optimizer='adam',  # Adam优化器（自适应学习率）
    loss='sparse_categorical_crossentropy',  # 稀疏交叉熵损失（标签为整数，非独热编码）
    metrics=['accuracy']  # 评估指标：准确率
)

# ==================== 步骤4：训练模型 ====================
history = model.fit(
    x_train, y_train,
    batch_size=128,  # 批次大小
    epochs=10,       # 训练轮数
    validation_split=0.2  # 从训练集划分20%作为验证集
)

# ==================== 步骤5：评估模型 ====================
test_loss, test_acc = model.evaluate(x_test, y_test)
print(f"\n测试集准确率: {test_acc:.4f}")

# ==================== 步骤6：预测演示（手写数字识别+加法计算） ====================
def predict_and_calculate(index1, index2):
    """预测两个手写数字并计算它们的和"""
    # 获取测试集样本
    sample1 = x_test[index1].reshape(1, -1)  # 调整为模型输入形状（1,784）
    sample2 = x_test[index2].reshape(1, -1)
    
    # 预测数字
    pred1 = np.argmax(model.predict(sample1))
    pred2 = np.argmax(model.predict(sample2))
    
    # 获取真实标签（可选对比）
    true1 = y_test[index1]
    true2 = y_test[index2]
    
    # 计算和
    sum_pred = pred1 + pred2
    
    # 可视化结果
    plt.figure(figsize=(10, 4))
    plt.subplot(1, 3, 1)
    plt.imshow(x_test[index1].reshape(28, 28), cmap='gray')
    plt.title(f"样本1\n预测: {pred1}, 真实: {true1}")
    
    plt.subplot(1, 3, 2)
    plt.imshow(x_test[index2].reshape(28, 28), cmap='gray')
    plt.title(f"样本2\n预测: {pred2}, 真实: {true2}")
    
    plt.subplot(1, 3, 3)
    plt.text(0.5, 0.5, f"预测和: {sum_pred}", 
             fontsize=24, ha='center', va='center')
    plt.axis('off')
    plt.suptitle("手写数字识别 + 加法计算演示", fontsize=16)
    plt.show()
    
    return pred1, pred2, sum_pred

# 测试：随机选择两个测试样本（可修改index1和index2指定具体样本）
index1 = 0  # 第一个测试样本索引
index2 = 1  # 第二个测试样本索引
pred1, pred2, sum_result = predict_and_calculate(index1, index2)
print(f"\n预测结果：{pred1} + {pred2} = {sum_result}")